import torch
import torch.nn as nn
import torch.nn.functional as F
from model.init_weight import init_weight
def get_sizes_list(dim, chunks):
    split_size = (dim + chunks - 1) // chunks
    sizes_list = [split_size] * chunks
    sizes_list[-1] = sizes_list[-1] - (sum(sizes_list) - dim) # Adjust last
    assert sum(sizes_list) == dim
    if sizes_list[-1]<0:
        n_miss = sizes_list[-2] - sizes_list[-1]
        sizes_list[-1] = sizes_list[-2]
        for j in range(n_miss):
            sizes_list[-j-1] -= 1
        assert sum(sizes_list) == dim
        assert min(sizes_list) > 0
    return sizes_list

def get_chunks(x,sizes):
    out = []
    begin = 0
    for s in sizes:
        y = x.narrow(1,begin,s)
        out.append(y)
        begin += s
    return out


class MFB(nn.Module): # bs, 20, 256  它是一个用于特征融合的模型，用于将两个输入特征向量进行乘法融合，并输出一个融合后的特征向量。

    def __init__(self,
            input_dims, #[256,256],256
            output_dim,
            mm_dim=256,
            factor=2,
            activ_input='elu', # change to elu
            activ_output='elu',
            normalize=False,
            dropout_input=0.,      #包括三个线性层（linear0、linear1和linear_out），用于特征的线性变换；
                                   # 还包括一些其他参数，如输入维度(input_dims)、输出维度(output_dim)、激活函数(activ_input和activ_output)、是否进行归一化(normalize)
            dropout_pre_norm=0.,
            dropout_output=0.):
        super(MFB, self).__init__()
        self.input_dims = input_dims
        self.mm_dim = mm_dim
        self.factor = factor
        self.output_dim = output_dim
        self.activ_input = activ_input
        self.activ_output = activ_output
        self.normalize = normalize
        self.dropout_input = dropout_input
        self.dropout_pre_norm = dropout_pre_norm
        self.dropout_output = dropout_output
        # Modules
        self.linear0 = nn.Linear(input_dims[0], mm_dim*factor) #256, 256
        self.linear1 = nn.Linear(input_dims[1], mm_dim*factor)
        self.linear_out = nn.Linear(mm_dim, output_dim)
        self.n_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
         #一个用于计算模型中可训练参数数量的属性
        self.init_all_weight()

    def init_all_weight(self):#对模型的所有权重进行初始化，调用了init_weight函数。
        self.linear0.apply(init_weight)
        self.linear1.apply(init_weight)
        self.linear_out.apply(init_weight)

    def forward(self, x): # bs, 20, 256
        x0 = self.linear0(x[0]) # Feature1   
        x1 = self.linear1(x[1]) # Feature2
        #通过两个线性层分别对输入特征进行线性变换得到x0和x1。然后根据activ_input指定的激活函数对x0和x1进行激活操作（默认为ELU）。
        if self.activ_input: # relu
            x0 = getattr(F, self.activ_input)(x0) #在进行一个ReLU
            x1 = getattr(F, self.activ_input)(x1)
            #从PyTorch的F模块（通常指代torch.nn.functional）中获取指定名称的函数的方法。
           #如果self.activ_input的值为'relu'，那么getattr(F, self.activ_input)将返回F.relu函数。
        if self.dropout_input > 0:
            x0 = F.dropout(x0, p=self.dropout_input, training=self.training)
            x1 = F.dropout(x1, p=self.dropout_input, training=self.training)
        #如果指定了dropout_input大于0，则在x0和x1上进行输入层的dropout操作。     
        z = x0 * x1   #这个就是元素级别的乘法

        if self.dropout_pre_norm > 0:
            z = F.dropout(z, p=self.dropout_pre_norm, training=self.training)
          #如果指定了dropout_pre_norm大于0，则在z上进行归一化前的dropout操作。
        if len(z.shape)==2:
            z = z.view(z.size(0), self.mm_dim, self.factor)
        else:
            z = z.view(z.size(0), z.size(1), self.mm_dim, self.factor) #做一个128*2的分割？
        z = z.sum(-1)
        #根据z的形状进行维度变换，将其转换为合适的形状。然后对z进行求和操作，并进行归一化（如果normalize为True）。
        if self.normalize:
            z = torch.sqrt(F.relu(z)) - torch.sqrt(F.relu(-z))
            z = F.normalize(z, p=2)

        z = self.linear_out(z)
#使用线性层linear_out对z进行线性变换并输出结果。最后，根据activ_output指定的激活函数对输出进
# 行激活操作（默认为ELU）。如果指定了dropout_output大于0，则在输出上进行输出层的dropout操作。最终，返回融合后的特征向量z。
        if self.activ_output:
            z = getattr(F, self.activ_output)(z)

        if self.dropout_output > 0:
            z = F.dropout(z, p=self.dropout_output, training=self.training)
        return z

